from dl_modules import Tokenizer, Attemtion, MHP_RNN

import torch

tokenizer = Tokenizer(3,4)
input_data = torch.randn(5,3)
torch.onnx.export(tokenizer,input_data,'onnx/Tokenizer.onnx',
    training=torch.onnx.TrainingMode.EVAL,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input':{0:'batch'},
        'output':{0:'batch'}
    })

attemtion = Attemtion(6,9,3)
input_data = torch.randn(5,8,6)
torch.onnx.export(attemtion,(input_data,input_data,input_data),'onnx/Attemtion.onnx',
    training=torch.onnx.TrainingMode.EVAL,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input':{0:'batch'},
        'output':{0:'batch'}
    })

mHP_RNN = MHP_RNN(6,9,3)
input_data = torch.randn(5,6)
torch.onnx.export(mHP_RNN,input_data,'onnx/MHP_RNN.onnx',
    training=torch.onnx.TrainingMode.EVAL,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input':{0:'batch'},
        'output':{0:'batch'}
    })